%%% Written by  Daniel Lewis, 2020, for "Robust Inference in Models 
%%% Identified via Heteroskedasticity", Review of Economics and Statistics.
%
% This file conducts the simulation study underlying Figure 2 and Tables 2
% and 3 in the paper and produces intermediate results (size-adjusted 
% critical values) for the power curves in Figure 5.

clear; clc
warning('off','all');
options=optimset('Display','off');
%% Setup

% Setting up parameters
T_emp=833;
Tc_emp=759;
params=[.7035 -.3098 .0039 0.0001 0.0070 0.0004]; % parameter values from empirical application (specification 1)
factor=params(5)/params(3); % proportional variance change for non-policy shock
dtilde_T=params(6)-factor*params(4); % equivalent to sigma^2_{2,c}*sigma^2_{1,p}/sigma^2_{1,c}*d/sqrt(T)
dtilde=dtilde_T*sqrt(T_emp); % equivalent to sigma^2_{2,c}*sigma^2_{1,p}/sigma^2_{1,c}*d - since we care about sigma^2_{2,p} as a function of d, not d itself, this is all we need
Tfrac=Tc_emp/T_emp; % fraction of control days in empirical sample
Trange=[400,800,1600]; % range of T values for simulation
ident=[1/10,1,10]*dtilde; % degrees of identification for simulation
N=10000; % Number of Monte Carlo samples
H=[1,params(2);params(1),1]; % Empirical H matrix
V1=diag(params(3:4)); % Structural variances for control regime

% Storage
Fstat=NaN(N,length(ident),3);
Sstat=NaN(N,length(ident),3);
Fstatp=NaN(N,length(ident),3);
Sstatp=NaN(N,length(ident),3);
Table2=zeros(length(ident),6);
Table3=zeros(length(ident),9);
quanttab=zeros(length(ident),3);
tstore=NaN(N,length(ident),3);

% Compute symbolic Jacobian
Hs=sym('H',2); % symbolic H matrix
assume(Hs,'real') ;
Hs(1)=1;Hs(4)=1; % unit diagonal normalization
Vc=diag(sym('Vc', [2 1])); % control regime structural variances
Vp=diag(sym('Vp', [2 1])); % policy regime structural variances
moms=[reshape(Hs*Vc*Hs',4,1);reshape(Hs*Vp*Hs',4,1)]; % compute moments
moms([3 7])=[]; % eliminate redundant moments
theta=[Hs(2); Hs(3); diag(Vc); diag(Vp)]; % stacked parameter vector
D=jacobian(moms,theta); % compute symbolic Jacobian

rng(240216); % seed random number generator

%% Runs

for j=1:length(ident) % looping over identification strength
    
    fprintf('Identification strength: %d\n',j);
    drun=ident(j); % set identification strength
    
    for t=1:length(Trange) % looping over sample lengths
        
        % setting up parameters and Jacobian
        T=Trange(t);
        fprintf('Sample length: %d\n',T);
        V2=[params(5),0;0,factor*params(4)+drun/sqrt(T)]; % scaling variance by sqrt(T) according with local-to-unity device
        truen=[params(1:5),V2(end)]'; % set true parameter values for this calibration - note this depends on drun
        Tc=floor(Tfrac*T); % number of control observations
        Tp=T-Tc; % number of policy observations
        Dt=D; % rescaling rows of symbolic Jacobian by sample share
        Dt(1:3,:)=D(1:3,:)*Tc/T;
        Dt(4:6,:)=D(4:6,:)*Tp/T;
        grad=matlabFunction(Dt,'Vars',theta); % convert symbolic Jacobian to function for fast evaluation
        
        parfor i=1:N
            rng(i); % seed rng
            % storage
            etasq1=NaN(Tc,3);
            etasq2=NaN(Tp,3);
            
            % draw data
            eps1=mvnrnd([0;0],V1,Tc)'; % shocks in control regime
            eps2=mvnrnd([0;0],V2,T-Tc)'; % shocks in policy regime
            eta1=H*eps1; % compute innovations
            eta2=H*eps2; % compute innovations
            %             eta1=eta1-mean([eta1,eta2],2); % demean so outer product equals covariance
            %             eta2=eta2-mean([eta1,eta2],2); % demean so outer product equals covariance
            eta1=eta1-mean(eta1,2); % demean so outer product equals covariance
            eta2=eta2-mean(eta2,2); % demean so outer product equals covariance
            
            % estimate structural parameters
            for s=1:length(eta1) % compute outer product of innovations in control regime
                mat=eta1(:,s)*eta1(:,s)';
                etasq1(s,:)=[mat(1:2),mat(4)]; % store vech
            end
            for s=1:length(eta2) % compute outer product of innovations in policy regime
                mat=eta2(:,s)*eta2(:,s)';
                etasq2(s,:)=[mat(1:2),mat(4)]; % store vech
            end
            cov1=eta1*eta1'/Tc; % covariance in each regime
            cov2=eta2*eta2'/Tp;
            mat=cov2*cov1^-1; % ratio of covariance matrices
            [a,b]=eig(mat);  % estimate columns of H using Brunnermeier et al (2019) approach
            if b(1)>b(4) % if the first shock has the largest proportional variance change
                Hint=a(1,1)/a(2); % parameter of interest is 1,1 element, but need to re-order and impose unit-diagonal normalization
                Ho=a(2,2)/a(3); % doing the same for other free element in H
            else
                Hint=a(1,2)/a(4); % parameter of interest is 1,2 element, need to impose unit-diagonal normalization
                Ho=a(2,1)/a(1); % doing the same for other free element in H
            end
            Hhat=[1,Hint;Ho,1]; % estimated H matrix, correctly ordered and normalized
            n1=Hhat^-1*cov1*Hhat'^-1; % structural variance matrix in control regime
            n2=Hhat^-1*cov2*Hhat'^-1; % structural variance matrix in policy regime
            est=[Ho,Hint,n1(1),n1(4),n2(1),n2(4)]; % putting together estimated parameter vector
            
            % compute test statistics
            [~,~,omega]=momfast(etasq1,etasq2,T,est); % compute moment covariance matrix
            ghat=grad(Ho,Hint,n1(1),n1(4),n2(1),n2(4)); % evaluate Jacobian at estimated parameter values
            Vinv=ghat'*omega^-1*ghat; % inverse of asymptotic variance
            V=Vinv^-1; % compute asymptotic variance
            tstore(i,j,t)=sqrt(T)*(est(2)-truen(2))/sqrt(V(2,2)); % compute and store t-statistic for H_12
            Fstat(i,j,t)=T*(est-truen')*Vinv*(est-truen')'/6; % compute and store F-statistic for full parameter vector
            Fstatp(i,j,t)=T*(est(2)-truen(2))'*V(2,2)^-1*(est(2)-truen(2))/1; % compute and store F-statistic for H_12
            Sstat(i,j,t)=momfast(etasq1,etasq2,T,truen); % compute S-statistic (= K statistic b/c just-identified) at estimates
            [est,S,flag]=fmincon(@(p) momfastsub(etasq1,etasq2,T,p,H(1,2),2),truen([1 3:end]),[],[],[],[],[-Inf; eps*ones(4,1)],[],@ordering,options); % minimize objective function conditional on null value of H_12
            if flag>0
                Sstatp(i,j,t)=S;
            end
            if 1000*floor(i/1000)==i % report progress every 1000 samples
                fprintf('%d draws complete\n',i);
            end
        end
        
    end
end
%% Figure 2
figure
labs={'$\hat{d}/10$','$\hat{d}$','10$\times\hat{d}$'}; % labels for strength of identification
subplot1(3,3,'FontS',12,'Gap',[.01,.01]); % set up subplots
xpt=-10:.01:10; % grid of x for reference normal pdf curve
ypt=normpdf(xpt); % y values for reference normal pdf curve
for t=1:3 % looping over sample sizes
    for j=1:3 % looping over indentification strengths
        subplot1(3*(t-1)+j)
        tstoreplot=squeeze(tstore(:,j,t)); % select data "slice"
        tstoreplot(tstoreplot<-1000)=[]; % drop any extreme outliers to keep axes readable
        tstoreplot(tstoreplot>1000)=[];
        histogram(real(tstoreplot),'BinWidth',.1,'EdgeColor','none','Normalization','pdf');xlim([-10,10]);
        hold on
        plot(xpt,ypt,'k','LineWidth',1.5) % plot reference normal pdf
        hold off
        yticks([.5 1])
        if 3*(t-1)+j==4||3*(t-1)+j==7||3*(t-1)+j==1
            ylabel(strcat('T=',num2str(Trange(t)))); % labels on the left margin
        end
        if t==3
            xlabel(char(labs(j)),'Interpreter','LaTex') % labels on the bottom margin
        end
        ylim([0,1]) % set ylimits
        xlim([-3,10]) % set xlimits
    end
end
subplot1(3)
legend('Estimated t-statistics','Normal pdf')


%% Tables

% Table 2: full parameter vector
for t=1:3
    for j=1:length(ident)
        F=squeeze(Fstat(:,j,t));
        F(isnan(F)) = []; % drop any NaNs
        S=squeeze(Sstat(:,j,t));
        S(isnan(S)) = []; % drop any NaNs
        Table2(t,2*(j-1)+1:2*j)=[mean(F>finv(.95,6,inf)),mean(S>chi2inv(.95,6))]; % compute rejection rates
    end
end
Table2

% Table 3: subset tests
for t=1:length(Trange)
    for j=1:length(ident)
        F=squeeze(Fstatp(:,j,t));
        F(isnan(F)) = []; % drop any NaNs
        S=squeeze(Sstatp(:,j,t));
        S(isnan(S)) = []; % drop any NaNs
        Table3(t,3*(j-1)+1:3*j)=[mean(F>finv(.95,1,inf)),mean(S>chi2inv(.95,6)),mean(S>chi2inv(.95,1))]; % compute rejection rates
    end
end
Table3

% Table of quantiles for size-adjusted power
for t=1:length(Trange)
    for j=1:length(ident)
        quanttab(j,2*(t-1)+1:2*t)=[quantile(squeeze(Fstatp(:,j,t)),.95),quantile(squeeze(Sstatp(:,j,t)),.95)]; % 95th percentile of each subset test statistic
    end
end
save('quanttab','quanttab')

%% Shock ordering constraint
function [c,ceq] = ordering(thetahat) % imposes ordering constraint: largest proportional variance change is in policy shock, epsilon 2
c = thetahat(4)/thetahat(2)-thetahat(5)/thetahat(3);
ceq = [];
end